from langchain_community.document_loaders.csv_loader import CSVLoader
from pydantic import BaseModel


class CustomArgs(BaseModel):
    delimiter: str = ","
    quotechar: str = '"'
    fieldnames: list[str] = None  # 自定义列名，默认会把首行作为列名
    source_column: str = None  # 自动义源，默认为None


async def load(path: str, custom_args: CustomArgs) -> list:
    loader = CSVLoader(file_path=path,
                       source_column=custom_args.source_column,
                       csv_args={
                           "delimiter": custom_args.delimiter,
                           "quotechar": custom_args.quotechar,
                           "fieldnames": custom_args.fieldnames
                       })
    data = loader.load()
    return data
